{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 解密 DeepSeek 之  \n",
    "## 完全从零实现的 MLA 算法\n",
    "### Multi-Head Latent Attention\n",
    "- 运行代码完全手写，一行一行带着运行，边写边知识点\n",
    "- 欢迎阅读**chaofa用代码打点酱油**的\n",
    "   - [博客原文](https://bruceyuan.com/post/hands-on-deepseek-mla.html)\n",
    "   - 实现稍有不一样，但原理是一样的，原因是一个是录制的过程中会有一些命名的改变\n",
    "- 观看本视频之前，需要有一定的基础知识\n",
    "   - 了解自注意力，self-attention\n",
    "   - 了解 multi-head self-attention\n",
    "   - 了解 Group Query Attention\n",
    "   - 了解 kv cache 是什么"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 基础的 package\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from typing import Optional, Tuple\n",
    "import math\n",
    "from dataclasses import dataclass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 一些前置代码，本次课暂时不涉及；\n",
    "# 如果有需要，以后可以专门出视频讲解\n",
    "class DeepseekV2RMSNorm(nn.Module):\n",
    "    def __init__(self, hidden_size, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.ones(hidden_size))\n",
    "        self.variance_epsilon = eps\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        input_dtype = hidden_states.dtype\n",
    "        hidden_states = hidden_states.to(torch.float32)\n",
    "        variance = hidden_states.pow(2).mean(-1, keepdim=True)\n",
    "        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)\n",
    "        return self.weight * hidden_states.to(input_dtype)\n",
    "    \n",
    "class DeepseekV2RotaryEmbedding(nn.Module):\n",
    "    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):\n",
    "        super().__init__()\n",
    "\n",
    "        self.dim = dim\n",
    "        self.max_position_embeddings = max_position_embeddings\n",
    "        self.base = base\n",
    "        inv_freq = 1.0 / (\n",
    "            self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)\n",
    "        )\n",
    "        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n",
    "        # 较小索引位置对应较低频率\n",
    "        # 较大的索引位置有较高的频率\n",
    "        \n",
    "        # Build here to make `torch.jit.trace` work.\n",
    "        self._set_cos_sin_cache(\n",
    "            seq_len=max_position_embeddings,\n",
    "            device=self.inv_freq.device,\n",
    "            dtype=torch.get_default_dtype(),\n",
    "        )\n",
    "        self.max_seq_len_cached = None\n",
    "\n",
    "    def _set_cos_sin_cache(self, seq_len, device, dtype):\n",
    "        self.max_seq_len_cached = seq_len\n",
    "        t = torch.arange(\n",
    "            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype\n",
    "        )\n",
    "\n",
    "        freqs = torch.outer(t, self.inv_freq.to(t.device))\n",
    "        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n",
    "        emb = torch.cat((freqs, freqs), dim=-1)\n",
    "        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n",
    "        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n",
    "\n",
    "    def forward(self, x, seq_len=None):\n",
    "        # x: [bs, num_attention_heads, seq_len, head_size]\n",
    "        if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:\n",
    "            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)\n",
    "\n",
    "        return (\n",
    "            self.cos_cached[:seq_len].to(dtype=x.dtype),\n",
    "            self.sin_cached[:seq_len].to(dtype=x.dtype),\n",
    "        )\n",
    "\n",
    "# Copied from transformers.models.llama.modeling_llama.rotate_half\n",
    "def rotate_half(x):\n",
    "    \"\"\"Rotates half the hidden dims of the input.\"\"\"\n",
    "    x1 = x[..., : x.shape[-1] // 2]\n",
    "    x2 = x[..., x.shape[-1] // 2 :]\n",
    "    return torch.cat((-x2, x1), dim=-1)\n",
    "\n",
    "# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb\n",
    "def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):\n",
    "    cos = cos[position_ids].unsqueeze(unsqueeze_dim)\n",
    "    sin = sin[position_ids].unsqueeze(unsqueeze_dim)\n",
    "\n",
    "    b, h, s, d = q.shape\n",
    "    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)\n",
    "\n",
    "    b, h, s, d = k.shape\n",
    "    k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)\n",
    "\n",
    "    q_embed = (q * cos) + (rotate_half(q) * sin)\n",
    "    k_embed = (k * cos) + (rotate_half(k) * sin)\n",
    "    return q_embed, k_embed\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 16, 1024, 192]) torch.Size([2, 16, 1024, 192])\n",
      "torch.Size([2, 1024, 7168])\n",
      "torch.Size([2, 16, 1024, 1024])\n"
     ]
    }
   ],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DeepseekConfig:\n",
    "    hidden_size: int\n",
    "    num_heads: int\n",
    "    max_position_embeddings: int  # 这是rope 相关的参数\n",
    "    rope_theta: float  # 频率，一般设置的比较大\n",
    "    \n",
    "    attention_dropout: float\n",
    "\n",
    "    q_lora_rank: int  # latent 的shape , 一般设置比较大一点；10000+\n",
    "    qk_rope_head_dim: int  # 64 \n",
    "    kv_lora_rank: int # 公式 41, 可能是 512\n",
    "\n",
    "    v_head_dim: int  # 128\n",
    "    qk_nope_head_dim: int\n",
    "    attention_bias: bool\n",
    "\n",
    "\n",
    "class MLA(nn.Module):\n",
    "    def __init__(self, config: DeepseekConfig):\n",
    "        super().__init__()\n",
    "        # 三个部分；\n",
    "        # part1 , mha 部分\n",
    "\n",
    "        self.attention_dropout = config.attention_dropout\n",
    "        self.hidden_size = config.hidden_size\n",
    "        self.num_heads = config.num_heads\n",
    "        self.v_head_dim = config.v_head_dim\n",
    "\n",
    "        self.out_proj = nn.Linear(\n",
    "            self.num_heads * self.v_head_dim,\n",
    "            self.hidden_size,\n",
    "            bias=False, # 可以加，也可以不加\n",
    "        )\n",
    "\n",
    "        # 最重要的是 part2\n",
    "        # MLA 压缩部分\n",
    "        # p2.1 down 压缩\n",
    "        self.qk_nope_head_dim = config.qk_nope_head_dim\n",
    "        self.qk_rope_head_dim = config.qk_rope_head_dim\n",
    "\n",
    "        self.q_lora_rank = config.q_lora_rank\n",
    "        # 一般会从 7168 -> 1536； 压缩比是 1/4.7\n",
    "\n",
    "        self.kv_lora_rank = config.kv_lora_rank\n",
    "        # 包含两个部分\n",
    "\n",
    "        self.q_down_proj = nn.Linear(\n",
    "            self.hidden_size,\n",
    "            self.q_lora_rank,\n",
    "            bias=config.attention_bias,\n",
    "        )\n",
    "        self.q_down_norm = DeepseekV2RMSNorm(self.q_lora_rank)\n",
    "\n",
    "        self.kv_down_proj = nn.Linear(\n",
    "            self.hidden_size,\n",
    "            self.kv_lora_rank + config.qk_rope_head_dim,\n",
    "            bias=config.attention_bias,\n",
    "        ) # qk_rope_head_dim 这个一般设置的很小，一般是 64\n",
    "        self.kv_down_norm = DeepseekV2RMSNorm(self.kv_lora_rank)\n",
    "        # down 之后包含了两个部分，要做 split \n",
    "\n",
    "        # p2.2 升维\n",
    "        # q and k shape is same\n",
    "        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim\n",
    "        self.q_up_proj = nn.Linear(\n",
    "            self.q_lora_rank,\n",
    "            self.num_heads * self.q_head_dim,\n",
    "            bias=config.attention_bias,\n",
    "        ) # 这里也要 split \n",
    "\n",
    "        self.kv_up_proj = nn.Linear(\n",
    "            self.kv_lora_rank,\n",
    "            self.num_heads * (\n",
    "                self.q_head_dim - config.qk_rope_head_dim + self.v_head_dim\n",
    "            ), # self.q_head_dim - config.qk_rope_head_dim = nope_shape ,\n",
    "            bias=config.attention_bias,\n",
    "        )\n",
    "\n",
    "        # part3: rope部分\n",
    "        self.rotary_emb = DeepseekV2RotaryEmbedding(\n",
    "            config.qk_rope_head_dim,\n",
    "            config.max_position_embeddings,\n",
    "            config.rope_theta,\n",
    "        ) \n",
    "\n",
    "        # part4: kv cache 的实现，下次课来讲，本次课先不管\n",
    "    \n",
    "    def forward(self, hidden_states, position_ids, attention_mask=None, ):\n",
    "        # hidden_states (b, seq_len, hidden_dim)\n",
    "        bsz, q_len, _ = hidden_states.size()\n",
    "\n",
    "        # 1. q compression\n",
    "        q = self.q_down_proj(\n",
    "            hidden_states\n",
    "        )\n",
    "        q = self.q_down_norm(q)\n",
    "        q = self.q_up_proj(q)\n",
    "        # q shape 是什么：self.num_heads * self.q_head_dim,\n",
    "        # (b, seq_len, self.num_heads * self.q_head_dim,)\n",
    "        q = q.view(\n",
    "            bsz, q_len, self.num_heads, self.q_head_dim\n",
    "        ).transpose(1, 2) \n",
    "        # （b, num_head, seq_len, q_head_dim）\n",
    "\n",
    "        q_nope, q_rope = torch.split(\n",
    "            q, \n",
    "            [self.qk_nope_head_dim, self.qk_rope_head_dim],\n",
    "            dim=-1\n",
    "        )\n",
    "\n",
    "        # kv part \n",
    "        c_kv = self.kv_down_proj(hidden_states)\n",
    "        c_kv, k_rope = torch.split(\n",
    "            c_kv,\n",
    "            [self.kv_lora_rank, self.qk_rope_head_dim],  # self.kv_lora_rank + config.qk_rope_head_dim\n",
    "            dim=-1,\n",
    "        ) # k_rope 的 shape （b, seq_len, elf.qk_rope_head_dim）\n",
    "        k_rope = k_rope.view(\n",
    "            bsz, q_len, 1, self.qk_rope_head_dim,\n",
    "        ).transpose(1, 2)  # boradcast\n",
    "        # (b, 1, seq_len, qk_rope_head_dim）)\n",
    "\n",
    "        kv = self.kv_down_norm(c_kv)\n",
    "        kv = self.kv_up_proj(kv)\n",
    "        # （b, seq, num_head * (\n",
    "            #     self.q_head_dim - config.qk_rope_head_dim + self.v_head_dim,\n",
    "            # )\n",
    "\n",
    "        kv = kv.view(\n",
    "            bsz, q_len, self.num_heads, \n",
    "            self.qk_nope_head_dim + self.v_head_dim\n",
    "        ).transpose(1, 2)\n",
    "\n",
    "        k_nope, value_states = torch.split(\n",
    "            kv,\n",
    "            [self.qk_nope_head_dim, self.v_head_dim],\n",
    "            dim=-1,\n",
    "        )\n",
    "\n",
    "        # apply 位置编码 rope\n",
    "        kv_seq_len = value_states.shape[-2]\n",
    "        # value_states shape (b, nums_head, seq_len, v_head_dim)\n",
    "\n",
    "        # 怎么使用 rope 下次课讲 ROPE 的时候来讲\n",
    "        cos, sin = self.rotary_emb(\n",
    "            value_states, seq_len=kv_seq_len,\n",
    "        )\n",
    "        q_rope, k_rope = apply_rotary_pos_emb(\n",
    "            q_rope, k_rope, cos, sin, position_ids,\n",
    "        )\n",
    "\n",
    "        # MHA\n",
    "        query_states = torch.concat(\n",
    "            [q_nope, q_rope], dim=-1\n",
    "        )\n",
    "        key_states = torch.concat(\n",
    "            [k_nope, k_rope.expand(-1, self.num_heads, -1, -1)], dim=-1\n",
    "        ) # (b, 1, seq_len, dim)\n",
    "        # shape is( b, num_head,\n",
    "        #  q_len, head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim)\n",
    "\n",
    "        print(query_states.shape, key_states.shape)\n",
    "\n",
    "        # MHA 无数遍了\n",
    "\n",
    "        attn_weights = torch.matmul(\n",
    "            query_states, key_states.transpose(2, 3)\n",
    "        )\n",
    "        attn_weights = attn_weights / math.sqrt(self.q_head_dim)\n",
    "\n",
    "        if attention_mask is not None:\n",
    "            # causal mask # \n",
    "            attn_weights = torch.masked_fill(\n",
    "                attn_weights,\n",
    "                attention_mask == 0,\n",
    "                float('-inf')\n",
    "            )\n",
    "        \n",
    "        # softmax 以及 output proj \n",
    "        attn_weights = F.softmax(\n",
    "            attn_weights, dim=-1\n",
    "        ).to(query_states.dtype)\n",
    "\n",
    "        attn_weights = F.dropout(\n",
    "            attn_weights, p=self.attention_dropout,\n",
    "            training=self.training,\n",
    "        )\n",
    "\n",
    "        output = torch.matmul(\n",
    "            # (b, num_head, q_len, q_len)\n",
    "            # value (b, nums_head, seq_len, v_head_dim)\n",
    "            attn_weights, value_states\n",
    "        )\n",
    "        output = output.transpose(1, 2).reshape(bsz, q_len, -1)\n",
    "\n",
    "        # （b, q_len, v_dim * num_head）\n",
    "\n",
    "        output = self.out_proj(\n",
    "            output\n",
    "        )\n",
    "\n",
    "        return output, attn_weights\n",
    "\n",
    "\n",
    "# 写一个测试函数\n",
    "def test_mla():\n",
    "    config = DeepseekConfig(\n",
    "        hidden_size=7168,\n",
    "        num_heads=16,\n",
    "        max_position_embeddings=1024,\n",
    "        rope_theta=128000,\n",
    "        attention_dropout=0.1,\n",
    "        q_lora_rank=1536,\n",
    "        qk_rope_head_dim=64,\n",
    "        kv_lora_rank=512,\n",
    "        \n",
    "        v_head_dim=128,\n",
    "        qk_nope_head_dim=128,\n",
    "        attention_bias=False,\n",
    "\n",
    "    )\n",
    "    \n",
    "    mla = MLA(config)\n",
    "    x = torch.randn(2, 1024, 7168)\n",
    "    position_ids = torch.arange(\n",
    "        config.max_position_embeddings,\n",
    "    ).unsqueeze(0).expand(\n",
    "        x.size(0), -1\n",
    "    ) # (batch_size, seq_len)\n",
    "    attn_output, attn_weights = mla(x, position_ids=position_ids)\n",
    "    print(attn_output.shape)\n",
    "    print(attn_weights.shape)\n",
    "\n",
    "\n",
    "test_mla()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
